import cv2
import numpy as np
from djitellopy import Tello  # Tello无人机控制库
import pyzbar.pyzbar as pyzbar  # 二维码识别库


class QRCodeController:
    def __init__(self):
        """初始化无人机连接和参数设置"""
        self.tello = Tello()
        self.tello.connect()  # 连接无人机
        self.tello.streamon()  # 开启视频流

        # 控制参数
        self.center_threshold = 30  # 中心点对齐阈值（像素单位）
        self.angle_threshold = 5  # 垂直角度阈值（度）
        self.Kp_yaw = 0.2  # 偏航（旋转）控制的比例系数
        self.Kp_vert = 0.1  # 垂直方向控制的比例系数

        # 状态变量
        self.qr_center = None  # 二维码中心坐标
        self.qr_angle = 0  # 二维码倾斜角度
        self.frame_center = (0, 0)  # 画面中心坐标
        self.last_qr_data = None  # 上一次识别的二维码内容（用于去重）

    def calculate_alignment(self, img):
        """
        检测二维码并计算对齐参数
        参数:
            img: 输入图像帧
        返回:
            二维码中心坐标, 多边形顶点, 解码数据
        """
        decoded = pyzbar.decode(img)
        if decoded:
            obj = decoded[0]  # 取第一个检测到的二维码
            # 计算二维码中心点坐标
            qr_center = (
                obj.rect.left + obj.rect.width // 2,
                obj.rect.top + obj.rect.height // 2
            )
            # 通过多边形顶点计算倾斜角度
            polygon = np.array(obj.polygon)
            vector = polygon[1] - polygon[0]  # 计算第一条边的向量
            self.qr_angle = np.degrees(np.arctan2(vector[1], vector[0]))
            return qr_center, polygon, obj.data.decode('utf-8')
        return None, None, None

    def draw_indicators(self, img, qr_center, polygon):
        """
        绘制视觉反馈元素
        参数:
            img: 要绘制的图像帧
            qr_center: 二维码中心坐标
            polygon: 二维码边界多边形顶点
        """
        # 绘制动态中心点（绿色表示已对齐，红色表示未对齐）
        center_color = (0, 255, 0) if self.is_centered() else (0, 0, 255)
        cv2.circle(img, self.frame_center, 8, center_color, -1)

        if qr_center:
            # 绘制二维码中心点（蓝色）
            cv2.circle(img, qr_center, 5, (255, 0, 0), -1)
            # 绘制画面中心到二维码中心的连线（青蓝色）
            cv2.line(img, self.frame_center, qr_center, (255, 255, 0), 2)

        if polygon is not None:
            # 绘制二维码边界（绿色表示垂直，红色表示倾斜）
            border_color = (0, 255, 0) if self.is_vertical() else (0, 0, 255)
            cv2.polylines(img, [polygon], True, border_color, 2)

    def is_centered(self):
        """判断二维码是否位于画面中心区域"""
        if self.qr_center is None:
            return False
        dx = self.qr_center[0] - self.frame_center[0]  # X轴偏移量
        dy = self.qr_center[1] - self.frame_center[1]  # Y轴偏移量
        # 计算欧氏距离并与阈值比较
        return np.sqrt(dx ** 2 + dy ** 2) < self.center_threshold

    def is_vertical(self):
        """判断二维码是否垂直（倾斜角度在阈值内）"""
        return abs(self.qr_angle) < self.angle_threshold

    def control_drone(self):
        """生成并发送无人机控制指令"""
        yaw_speed = 0  # 偏航（旋转）速度
        vert_speed = 0  # 垂直方向速度

        if self.qr_center:
            # 水平方向控制：计算X轴偏移量并应用比例控制
            x_offset = self.frame_center[0] - self.qr_center[0]
            yaw_speed = int(np.clip(x_offset * self.Kp_yaw, -100, 100))

            # 垂直方向控制：计算Y轴偏移量并应用比例控制
            y_offset = self.frame_center[1] - self.qr_center[1]
            vert_speed = int(np.clip(y_offset * self.Kp_vert, -100, 100))

        # 发送遥控指令（左右，前后，升降，旋转）
        self.tello.send_rc_control(
            0,  # 左右平移（此处未使用）
            vert_speed,  # 前后平移（实际用于控制高度）
            0,  # 升降速度（此处未使用）
            yaw_speed  # 旋转速度
        )


    def run(self):
        """主运行循环"""
        try:
            while True:
                # 获取视频流帧
                img = self.tello.get_frame_read().frame
                if img is None:
                    continue

                # 计算画面中心
                h, w = img.shape[:2]
                self.frame_center = (w // 2, h // 2)  # 更新画面中心坐标

                # 二维码检测与参数计算
                self.qr_center, polygon, qr_data = self.calculate_alignment(img)

                # 绘制视觉反馈元素
                self.draw_indicators(img, self.qr_center, polygon)

                # 二维码数据处理逻辑
                if self.is_centered() and self.is_vertical():
                    if qr_data and qr_data != self.last_qr_data:
                        print('识别到二维码数据:', qr_data)
                        self.last_qr_data = qr_data  # 记录已识别数据
                else:
                    self.last_qr_data = None  # 重置识别状态

                # 发送控制指令
                self.control_drone()

                # 显示视频流
                cv2.imshow("Tello Camera", cv2.cvtColor(img, cv2.COLOR_RGB2BGR))

                # 按Q键退出
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break

        finally:
            # 清理资源
            self.tello.streamoff()  # 关闭视频流
            self.tello.end()  # 断开连接
            cv2.destroyAllWindows()  # 关闭所有OpenCV窗口


if __name__ == "__main__":
    controller = QRCodeController()
    controller.run()